import sys
from collections import defaultdict

import numpy as np
import pygame
from pygame.locals import DOUBLEBUF, HWSURFACE

from driving_gym.data.data_provider import Snapshot


class Visualizer:
    def __init__(self, config: dict):
        self.config = config
        self.image_size_x = self.config.get("image_size_x", 168)
        self.image_size_y = self.config.get("image_size_y", 168)
        self.mark_sensor = self.config.get("mark_sensor", True)
        self.bg_color = self.config.get("background_color", (0, 0, 0))

        # Rendering
        self.screen = None
        self.font = None
        self.agent_sensors = defaultdict(list)
        self.total_agents = 0
        self.total_sensors_per_agent = 0
        self.window_width = 0
        self.window_height = 0
        self.initialized = False

    def parse_config(self, agent_configs: dict):
        """Prepares the visualizer by parsing agent configurations and calculating window size.

        Args:
            agent_configs (dict): Dictionary containing configurations for each agent.

        Example:
        {
            "hero": {"render": ["rgb", "bev"]},
            "npc": {"render": ["rgb"]}
        }
        """
        self.agent_sensors.clear()
        self.total_agents = len(agent_configs)
        max_sensors = 0

        for agent_name, agent_conf in agent_configs.items():
            sensors = agent_conf.get("render", [])
            self.agent_sensors[agent_name] = sensors
            if len(sensors) > max_sensors:
                max_sensors = len(sensors)

        self.total_sensors_per_agent = max_sensors

        # Calculate window size
        self.window_width = self.total_sensors_per_agent * self.image_size_x
        self.window_height = self.total_agents * self.image_size_y

    def render(self, snapshot: Snapshot):
        """Renders the sensor data onto the Pygame window.

        Args:
            snapshot (Snapshot): Snapshot data containing sensor images.
        """
        if not self.initialized:
            pygame.init()
            pygame.display.set_caption("Driving Gym")
            self.font = pygame.font.SysFont(None, 24)
            self.screen = pygame.display.set_mode(
                (self.window_width, self.window_height), DOUBLEBUF | HWSURFACE
            )
            self.initialized = True

        # Fill background
        self.screen.fill(self.bg_color)

        y_offset = 0
        for agent_idx, (agent_name, sensors) in enumerate(self.agent_sensors.items()):
            x_offset = 0
            sensor_datas = snapshot.data[agent_name]["sensors"]
            for sensor in sensors:
                if sensor in sensor_datas:
                    img_array = sensor_datas[sensor]
                else:
                    img_array = np.zeros(
                        (self.image_size_y, self.image_size_x, 3), dtype=np.uint8
                    )

                # Ensure the image has three channels
                if img_array.ndim == 2:
                    img_array = np.stack([img_array] * 3, axis=-1)
                elif img_array.shape[2] == 1:
                    img_array = np.concatenate([img_array] * 3, axis=2)

                # Resize image
                img_surface = pygame.surfarray.make_surface(
                    np.transpose(img_array, (1, 0, 2))
                )
                img_surface = pygame.transform.scale(
                    img_surface, (self.image_size_x, self.image_size_y)
                )

                # Blit image
                self.screen.blit(img_surface, (x_offset, y_offset))

                # Render sensor label if needed
                if self.mark_sensor:
                    label = f"{sensor}"
                    text_surface = self.font.render(label, True, (255, 255, 255))
                    self.screen.blit(text_surface, (x_offset + 5, y_offset + 5))

                x_offset += self.image_size_x

            y_offset += self.image_size_y

        pygame.display.flip()

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()

    def close(self):
        """Cleans up the Pygame resources."""
        if self.initialized:
            pygame.quit()
            self.initialized = False
